import os
import os.path as osp
import sys
import glfw
import glm
from OpenGL.GL import *
import numpy as np
import trimesh
import ctypes
import math
import copy
import torch
import smplx
import shutil

from PIL import Image
from learnopengl.ibl.shader import Shader
from scipy.spatial.transform import Rotation as R
from tqdm import tqdm
from multiprocessing import Pool, cpu_count


print_device = True
SCR_WIDTH = SCR_HEIGHT = 512
crop_scale = 1.05

euler_angle = [90, 0, 0]
r = R.from_euler('xyz', euler_angle, degrees=True)
rotation_matrix = r.as_matrix()
rotation_matrix_inv = np.linalg.inv(rotation_matrix)

# F:\workspaces\mmlab\cliff\data\body_models\smplx\models\smplx
# D:\workspaces\preprocessed_datasets\models_smplx_v1_1\models\smplx
m_smplx = smplx.SMPLX(model_path=r'F:\workspaces\mmlab\cliff\data\body_models\smplx\models\smplx', gender='neutral', ext='npz', use_pca=False)

# left = -5.0
# right = left + 10
# bottom = -0.2
# top = bottom + 10
# zNear = -100
# zFar = 100.0


# Vertex Shader Source
verticesShaderSource = '''
    #version 330 core
    layout (location = 0) in vec3 aPos;
    layout (location = 1) in vec2 aTexCoord;
    layout (location = 2) in vec3 aNormal;
    
    out vec2 TexCoord;
    out vec3 Normal;
    
    uniform mat4 model;
    uniform mat4 view;
    uniform mat4 projection;
    
    void main()
    {
        Normal = mat3(transpose(inverse(model))) * aNormal;
        gl_Position = projection * view * model * vec4(aPos, 1.0);
        TexCoord = aTexCoord;
    }
'''

# Fragment Shader Source
fragmentShaderSource = '''
    #version 330 core
    out vec4 FragColor;
    
    in vec2 TexCoord;

    uniform sampler2D ourTexture;
    
    void main()
    {
        FragColor = texture(ourTexture, TexCoord);
    }
'''

# Normal_Vert Shader Source
NormalVerticesShaderSource = '''
    #version 330 core
    layout (location = 0) in vec3 aPos;
    layout (location = 1) in vec3 aNormal;
    
    out vec3 Normal;
    
    uniform mat4 model;
    uniform mat4 view;
    uniform mat4 projection;
    
    void main()
    {
        Normal = mat3(transpose(inverse(model))) * aNormal;
        gl_Position = projection * view * model * vec4(aPos, 1.0);
    }
'''

# Normal Shader Source
normalShaderSource = '''
    #version 330 core
    out vec4 FragColor;
    
    in vec3 Normal;

    void main()
    {
        vec3 normals = normalize(Normal);
        vec3 color = normals * 0.5 + 0.5;
        FragColor = vec4(color, 1.0);
    }
'''

def load_texture(texture_path, texture_unit, flip=True):
    img = Image.open(texture_path).convert('RGB')
    if flip:
        img = img.transpose(Image.FLIP_TOP_BOTTOM)
    img_data = img.tobytes()
    width, height = img.size

    texture = glGenTextures(1)
    glActiveTexture(GL_TEXTURE0 + texture_unit)
    glPixelStorei(GL_UNPACK_ALIGNMENT, 1)
    glBindTexture(GL_TEXTURE_2D, texture)

    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, width, height, 0, GL_RGB, GL_UNSIGNED_BYTE, img_data)
    glGenerateMipmap(GL_TEXTURE_2D)
    
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAX_LEVEL, 3)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR_MIPMAP_LINEAR)

    return texture

def framebuffer_size_callback(window, width, height):
    glViewport(0, 0, width, height)

def processInput(window):
    if glfw.get_key(window, glfw.KEY_ESCAPE) == glfw.PRESS:
        glfw.set_window_should_close(window, True)

def save_normal_map(width, height, filepath, type_view):
    glReadBuffer(GL_COLOR_ATTACHMENT0)
    data = glReadPixels(0, 0, width, height, GL_RGBA, GL_UNSIGNED_BYTE)
    image = Image.frombytes("RGBA", (width, height), data)
    image = image.transpose(Image.FLIP_TOP_BOTTOM)
    if type_view == 'back':
        image = image.transpose(Image.FLIP_LEFT_RIGHT)
    image.save(filepath)

def save_render_map(width, height, filepath, type_view):
    glReadBuffer(GL_COLOR_ATTACHMENT0)
    data = glReadPixels(0, 0, width, height, GL_RGBA, GL_UNSIGNED_BYTE)
    image = Image.frombytes("RGBA", (width, height), data)
    image = image.transpose(Image.FLIP_TOP_BOTTOM)
    if type_view == 'back':
        image = image.transpose(Image.FLIP_LEFT_RIGHT)
    image.save(filepath)

# 将 3x3 旋转矩阵转换为 4x4 齐次矩阵
def to_homogeneous(matrix):
    assert matrix.shape == (3, 3)
    homogeneous_matrix = np.eye(4)
    homogeneous_matrix[:3, :3] = matrix
    return homogeneous_matrix

def normalize_vector(vertices, refer_vertices=None):
    if refer_vertices is not None:
        vertices_base = refer_vertices
    else:
        vertices_base = vertices
    min_vals = vertices_base.min(axis=0)
    max_vals = vertices_base.max(axis=0)
    center = (max_vals + min_vals) / 2.0
    scale = (max_vals - min_vals).max() / 2.0

    vector_new = (vertices - center) / scale
    return vector_new

def load_smplx(path):
    params = np.load(path, allow_pickle=True)

    smpl_out = m_smplx.forward(
        transl=torch.tensor(params['transl'], dtype=torch.float32),
        global_orient=torch.tensor(params['global_orient'], dtype=torch.float32),
        body_pose=torch.tensor(params['body_pose'], dtype=torch.float32),
        betas=torch.tensor(params['betas'], dtype=torch.float32),
        left_hand_pose=torch.tensor(params['left_hand_pose'], dtype=torch.float32),
        right_hand_pose=torch.tensor(params['right_hand_pose'], dtype=torch.float32),
        expression=torch.tensor(params['expression'], dtype=torch.float32),
        jaw_pose=torch.tensor(params['jaw_pose'], dtype=torch.float32),
        leye_pose=torch.tensor(params['leye_pose'], dtype=torch.float32),
        reye_pose=torch.tensor(params['reye_pose'], dtype=torch.float32)
    )
    
    smpl_joints = smpl_out.joints[0].detach()

    return smpl_joints


def calculate_bounding_box(vertices):
    min_coords = np.min(vertices, axis=0)
    max_coords = np.max(vertices, axis=0)
    return min_coords, max_coords


def process_clothed_files(mesh_file, smpl_file, texture_path, save_dir, SCR_WIDTH, SCR_HEIGHT, is_XRotate):
    # 加载网格和纹理
    mesh = trimesh.load(mesh_file, process=False)
    vertices_ori = np.array(mesh.vertices, dtype=np.float32)
    clothed_vertices = copy.deepcopy(vertices_ori)

    faces = np.array(mesh.faces.flatten(), dtype=np.int32)
    normals = np.array(mesh.vertex_normals, dtype=np.float32)
    uv_coords = np.array(mesh.visual.uv, dtype=np.float32)
    vertices = np.concatenate([vertices_ori, uv_coords, normals], axis=1)

    # 计算正交投影的坐标
    if is_XRotate:
        vertices_new = np.matmul(rotation_matrix_inv, clothed_vertices.T).T
    else:
        vertices_new = clothed_vertices
    min_x = vertices_new[:, 0].min()
    max_x = vertices_new[:, 0].max()
    min_y = vertices_new[:, 1].min()
    max_y = vertices_new[:, 1].max()
    
    # # left, right, bottom, top
    w = max_x - min_x
    h = max_y - min_y
    media_x = 0 #(max_x + min_x) / 2
    media_y = (max_y + min_y) / 2
    if w > h:
        w = w * crop_scale
        left = media_x - w / 2
        right = media_x + w / 2
        top = media_y + w / 2
        bottom = media_y - w / 2
    else:
        h = h * crop_scale
        left = media_x - h / 2
        right = media_x + h / 2
        top = media_y + h / 2
        bottom = media_y - h / 2
    zNear = -100
    zFar = 100.0

    # 初始化
    if not glfw.init():
        raise Exception("GLFW can't be initialized")
    glfw.window_hint(glfw.VISIBLE, glfw.FALSE)  # 隐藏窗口
    glfw.window_hint(glfw.CONTEXT_VERSION_MAJOR, 3)
    glfw.window_hint(glfw.CONTEXT_VERSION_MINOR, 3)
    glfw.window_hint(glfw.OPENGL_PROFILE, glfw.OPENGL_CORE_PROFILE)
    
    window = glfw.create_window(SCR_WIDTH, SCR_HEIGHT, "LearnOpenGL", None, None)
    if not window:
        glfw.terminate()
        raise Exception("GLFW window can't be created")

    glfw.make_context_current(window)
    glfw.set_framebuffer_size_callback(window, framebuffer_size_callback)
    
    glfw.set_input_mode(window, glfw.CURSOR, glfw.CURSOR_DISABLED)
    
    glEnable(GL_DEPTH_TEST)
    
    # 打印 GPU 信息
    global print_device
    if print_device:
        print("Renderer:", glGetString(GL_RENDERER).decode())
        print("OpenGL version supported:", glGetString(GL_VERSION).decode())
        print_device = False
    
    VAO = glGenVertexArrays(1)
    VBO = glGenBuffers(1)
    EBO = glGenBuffers(1)
    
    glBindVertexArray(VAO)
    
    glBindBuffer(GL_ARRAY_BUFFER, VBO)
    glBufferData(GL_ARRAY_BUFFER, vertices.nbytes, vertices, GL_STATIC_DRAW)
    
    glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, EBO)
    glBufferData(GL_ELEMENT_ARRAY_BUFFER, faces.nbytes, faces, GL_STATIC_DRAW)

    stride = vertices.itemsize * 8  # 3 for pos, 2 for tex, 3 for normal
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, stride, ctypes.c_void_p(0))
    glEnableVertexAttribArray(0)
    glVertexAttribPointer(1, 2, GL_FLOAT, GL_FALSE, stride, ctypes.c_void_p(3 * vertices.itemsize))
    glEnableVertexAttribArray(1)
    glVertexAttribPointer(2, 3, GL_FLOAT, GL_FALSE, stride, ctypes.c_void_p(5 * vertices.itemsize))
    glEnableVertexAttribArray(2)
    
    renderShaderProgram = Shader(verticesShaderSource, fragmentShaderSource)
    normalShaderProgram = Shader(verticesShaderSource, normalShaderSource)
    
    texture_map = load_texture(texture_path, 0)
    
    renderShaderProgram.use()
    glUniform1i(glGetUniformLocation(renderShaderProgram.ID, "ourTexture"), 0)

    # 设置帧缓冲对象
    FBO = glGenFramebuffers(1)
    glBindFramebuffer(GL_FRAMEBUFFER, FBO)
    
    # 创建纹理
    textureColorBuffer = glGenTextures(2)
    glBindTexture(GL_TEXTURE_2D, textureColorBuffer[0])
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, SCR_WIDTH, SCR_HEIGHT, 0, GL_RGBA, GL_UNSIGNED_BYTE, None)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
    glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, textureColorBuffer[0], 0)
    
    glBindTexture(GL_TEXTURE_2D, textureColorBuffer[1])
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, SCR_WIDTH, SCR_HEIGHT, 0, GL_RGBA, GL_UNSIGNED_BYTE, None)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
    glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT1, GL_TEXTURE_2D, textureColorBuffer[1], 0)
    
    # 创建渲染缓冲对象
    rbo = glGenRenderbuffers(1)
    glBindRenderbuffer(GL_RENDERBUFFER, rbo)
    glRenderbufferStorage(GL_RENDERBUFFER, GL_DEPTH24_STENCIL8, SCR_WIDTH, SCR_HEIGHT)
    glFramebufferRenderbuffer(GL_FRAMEBUFFER, GL_DEPTH_STENCIL_ATTACHMENT, GL_RENDERBUFFER, rbo)
    
    if glCheckFramebufferStatus(GL_FRAMEBUFFER) != GL_FRAMEBUFFER_COMPLETE:
        raise Exception("Framebuffer is not complete!")
    
    glBindFramebuffer(GL_FRAMEBUFFER, 0)
    
    y_rotate = np.arange(0, 360, 360 // 36)
    i = 0
    while not glfw.window_should_close(window):  
        processInput(window)
        
        model = glm.rotate(glm.mat4(1.0), math.radians(y_rotate[i]), glm.vec3(0, 1, 0))
        if is_XRotate:
            model = model @ glm.mat4(to_homogeneous(rotation_matrix_inv))
        projection = glm.ortho(left, right, bottom, top, zNear, zFar)
        view_front = glm.lookAt(glm.vec3(0, 0, 1), glm.vec3(0, 0, 0), glm.vec3(0, 1, 0))
        view_back = glm.lookAt(glm.vec3(0, 0, -1), glm.vec3(0, 0, 0), glm.vec3(0, 1, 0))

        def process_front_back(view, type_view="front"):
            glBindFramebuffer(GL_FRAMEBUFFER, FBO)
            glClearColor(0.0, 0.0, 0.0, 0.0)

            # # 渲染法线贴图
            glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
            normalShaderProgram.use()
            normalShaderProgram.set_mat4("projection", projection)
            normalShaderProgram.set_mat4("view", view)
            normalShaderProgram.set_mat4("model", model)
            
            glBindVertexArray(VAO)
            glDrawElements(GL_TRIANGLES, len(faces), GL_UNSIGNED_INT, None)
            os.makedirs(f'{save_dir}/view_{i:02d}', exist_ok=True)
            save_normal_map(SCR_WIDTH, SCR_HEIGHT, f'{save_dir}/view_{i:02d}/normal_{type_view}.png', type_view)

            # 渲染纹理贴图
            glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
            renderShaderProgram.use()
            renderShaderProgram.set_mat4("projection", projection)
            renderShaderProgram.set_mat4("view", view)
            renderShaderProgram.set_mat4("model", model)
            
            glActiveTexture(GL_TEXTURE0)
            glBindTexture(GL_TEXTURE_2D, texture_map)
            
            glBindVertexArray(VAO)
            glDrawElements(GL_TRIANGLES, len(faces), GL_UNSIGNED_INT, None)
            save_render_map(SCR_WIDTH, SCR_HEIGHT, f'{save_dir}/view_{i:02d}/image_{type_view}.png', type_view)

            glBindFramebuffer(GL_FRAMEBUFFER, 0)
        
        process_front_back(view_front, "front")
        process_front_back(view_back, "back")

        glfw.swap_buffers(window)
        glfw.poll_events()
        i += 1
        if i == len(y_rotate):
        # if i == 1:
            break
    
    glDeleteBuffers(1, [VBO])
    glDeleteBuffers(1, [EBO])
    glDeleteVertexArrays(1, [VAO])
    glDeleteTextures(1, textureColorBuffer)
    glDeleteRenderbuffers(1, [rbo])
    glDeleteFramebuffers(1, [FBO])
    glfw.terminate()


def process_smpl_files(mesh_file, smpl_file, save_dir, SCR_WIDTH, SCR_HEIGHT, is_XRotate):
    # 加载网格和纹理
    mesh = trimesh.load(smpl_file, process=False)
    vertices_ori = np.array(mesh.vertices, dtype=np.float32)

    clothed_mesh = trimesh.load(mesh_file, process=False)
    clothed_vertices = np.array(clothed_mesh.vertices, dtype=np.float32)

    faces = np.array(mesh.faces.flatten(), dtype=np.uint32)
    normals = np.array(mesh.vertex_normals, dtype=np.float32)
            
    vertices = np.concatenate([vertices_ori, normals], axis=1)
    
    # 计算正交投影的坐标
    if is_XRotate:
        vertices_new = np.matmul(rotation_matrix_inv, clothed_vertices.T).T
    else:
        vertices_new = clothed_vertices
    min_x = vertices_new[:, 0].min()
    max_x = vertices_new[:, 0].max()
    min_y = vertices_new[:, 1].min()
    max_y = vertices_new[:, 1].max()
    
    # # left, right, bottom, top
    w = max_x - min_x
    h = max_y - min_y
    media_x = 0 #(max_x + min_x) / 2
    media_y = (max_y + min_y) / 2
    if w > h:
        w = w * crop_scale
        left = media_x - w / 2
        right = media_x + w / 2
        top = media_y + w / 2
        bottom = media_y - w / 2
    else:
        h = h * crop_scale
        left = media_x - h / 2
        right = media_x + h / 2
        top = media_y + h / 2
        bottom = media_y - h / 2
    zNear = -100
    zFar = 100.0

    # 初始化
    if not glfw.init():
        raise Exception("GLFW can't be initialized")
    glfw.window_hint(glfw.VISIBLE, glfw.FALSE)  # 隐藏窗口
    glfw.window_hint(glfw.CONTEXT_VERSION_MAJOR, 3)
    glfw.window_hint(glfw.CONTEXT_VERSION_MINOR, 3)
    glfw.window_hint(glfw.OPENGL_PROFILE, glfw.OPENGL_CORE_PROFILE)
    
    window = glfw.create_window(SCR_WIDTH, SCR_HEIGHT, "LearnOpenGL", None, None)
    if not window:
        glfw.terminate()
        raise Exception("GLFW window can't be created")

    glfw.make_context_current(window)
    glfw.set_framebuffer_size_callback(window, framebuffer_size_callback)
    
    glfw.set_input_mode(window, glfw.CURSOR, glfw.CURSOR_DISABLED)
    
    glEnable(GL_DEPTH_TEST)
    
    # 打印 GPU 信息
    global print_device
    if print_device:
        print("Renderer:", glGetString(GL_RENDERER).decode())
        print("OpenGL version supported:", glGetString(GL_VERSION).decode())
        print_device = False
    
    VAO = glGenVertexArrays(1)
    VBO = glGenBuffers(1)
    EBO = glGenBuffers(1)
    
    glBindVertexArray(VAO)
    
    glBindBuffer(GL_ARRAY_BUFFER, VBO)
    glBufferData(GL_ARRAY_BUFFER, vertices.nbytes, vertices, GL_STATIC_DRAW)
    
    glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, EBO)
    glBufferData(GL_ELEMENT_ARRAY_BUFFER, faces.nbytes, faces, GL_STATIC_DRAW)

    stride = vertices.itemsize * 6  # 3 for pos, 3 for normal
    glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, stride, ctypes.c_void_p(0))
    glEnableVertexAttribArray(0)
    glVertexAttribPointer(1, 3, GL_FLOAT, GL_FALSE, stride, ctypes.c_void_p(3 * vertices.itemsize))
    glEnableVertexAttribArray(1)
    
    normalShaderProgram = Shader(NormalVerticesShaderSource, normalShaderSource)
    
    # 设置帧缓冲对象
    FBO = glGenFramebuffers(1)
    glBindFramebuffer(GL_FRAMEBUFFER, FBO)
    
    # 创建纹理
    textureColorBuffer = glGenTextures(2)
    glBindTexture(GL_TEXTURE_2D, textureColorBuffer[0])
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, SCR_WIDTH, SCR_HEIGHT, 0, GL_RGBA, GL_UNSIGNED_BYTE, None)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
    glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT0, GL_TEXTURE_2D, textureColorBuffer[0], 0)
    
    glBindTexture(GL_TEXTURE_2D, textureColorBuffer[1])
    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGBA, SCR_WIDTH, SCR_HEIGHT, 0, GL_RGBA, GL_UNSIGNED_BYTE, None)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
    glFramebufferTexture2D(GL_FRAMEBUFFER, GL_COLOR_ATTACHMENT1, GL_TEXTURE_2D, textureColorBuffer[1], 0)
    
    # 创建渲染缓冲对象
    rbo = glGenRenderbuffers(1)
    glBindRenderbuffer(GL_RENDERBUFFER, rbo)
    glRenderbufferStorage(GL_RENDERBUFFER, GL_DEPTH24_STENCIL8, SCR_WIDTH, SCR_HEIGHT)
    glFramebufferRenderbuffer(GL_FRAMEBUFFER, GL_DEPTH_STENCIL_ATTACHMENT, GL_RENDERBUFFER, rbo)
    
    if glCheckFramebufferStatus(GL_FRAMEBUFFER) != GL_FRAMEBUFFER_COMPLETE:
        raise Exception("Framebuffer is not complete!")
    
    glBindFramebuffer(GL_FRAMEBUFFER, 0)
    
    y_rotate = np.arange(0, 360, 360 // 36)
    i = 0
    while not glfw.window_should_close(window):  
        processInput(window)
        
        model = glm.rotate(glm.mat4(1.0), math.radians(y_rotate[i]), glm.vec3(0, 1, 0))
        if is_XRotate:
            model = model @ glm.mat4(to_homogeneous(rotation_matrix_inv))
        projection = glm.ortho(left, right, bottom, top, zNear, zFar)
        view_front = glm.lookAt(glm.vec3(0, 0, 1), glm.vec3(0, 0, 0), glm.vec3(0, 1, 0))
        view_back = glm.lookAt(glm.vec3(0, 0, -1), glm.vec3(0, 0, 0), glm.vec3(0, 1, 0))

        def process_front_back(view, type_view="front"):
            glBindFramebuffer(GL_FRAMEBUFFER, FBO)
            glClearColor(0.0, 0.0, 0.0, 0.0)

            # # 渲染法线贴图
            glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
            normalShaderProgram.use()
            normalShaderProgram.set_mat4("projection", projection)
            normalShaderProgram.set_mat4("view", view)
            normalShaderProgram.set_mat4("model", model)
            
            glBindVertexArray(VAO)
            glDrawElements(GL_TRIANGLES, len(faces), GL_UNSIGNED_INT, None)
            os.makedirs(f'{save_dir}/view_{i:02d}', exist_ok=True)
            save_normal_map(SCR_WIDTH, SCR_HEIGHT, f'{save_dir}/view_{i:02d}/T_normal_{type_view}.png', type_view)

            glBindFramebuffer(GL_FRAMEBUFFER, 0)
        
        process_front_back(view_front, "front")
        process_front_back(view_back, "back")

        glfw.swap_buffers(window)
        glfw.poll_events()
        i += 1
        if i == len(y_rotate):
        # if i == 1:
            break
    
    glDeleteBuffers(1, [VBO])
    glDeleteBuffers(1, [EBO])
    glDeleteVertexArrays(1, [VAO])
    glDeleteTextures(1, textureColorBuffer)
    glDeleteRenderbuffers(1, [rbo])
    glDeleteFramebuffers(1, [FBO])
    glfw.terminate()


def process_clothed_task(args):
    return process_clothed_files(*args)

def process_smpl_task(args):
    return process_smpl_files(*args)

def process_task(args):
    process_smpl_files(*args[1])
    process_clothed_files(*args[0])

if __name__ == "__main__":
    start = 532   #532
    rotate_i = 526
    src_dir = r"THuman2.1_Release"
    save_dir = r"THuman2.1_Release\image_data_v2"
    # if osp.exists(save_dir):
    #     shutil.rmtree(save_dir)
    input_dir = osp.join(src_dir, 'model')
    smpl_dir = osp.join(src_dir, 'smplx')

    filedir_list = os.listdir(input_dir)
    filedir_list.sort()

    file_groups = []
    for file_i in filedir_list:
        mesh_file = osp.join(input_dir, file_i, file_i+'.obj')
        smpl_file = osp.join(smpl_dir, file_i, 'mesh_smplx.obj')
        texture_path = osp.join(input_dir, file_i, 'material0.jpeg')
        assert osp.exists(mesh_file), mesh_file
        assert osp.exists(texture_path), texture_path
        file_groups.append({'mesh_file': mesh_file, 'smpl_file': smpl_file, 'texture_path': texture_path})

    os.makedirs(save_dir, exist_ok=True)
    exist_list = os.listdir(save_dir)
    exist_id = []
    for img in exist_list:
        exist_id.append(int(img))

    clothed_tasks = []
    smpl_task = []
    task = []
    for i, group in enumerate(file_groups):
        if i < start:
            continue
        if i < rotate_i:
            is_XRotate = False
        else:
            is_XRotate = True
        if i in exist_id:
            continue
        save_sub_dir = os.path.join(save_dir, os.path.basename(group['mesh_file']).split('.')[0])
        clothed_tasks.append((group['mesh_file'], group['smpl_file'], group['texture_path'], save_sub_dir, SCR_WIDTH, SCR_HEIGHT, is_XRotate))
        smpl_task.append((group['mesh_file'], group['smpl_file'], save_sub_dir, SCR_WIDTH, SCR_HEIGHT, is_XRotate))
        task.append([clothed_tasks[-1], smpl_task[-1]])
    

    
    num_workers = cpu_count()  # 获取系统中的 CPU 核心数
    print(f"Number of workers: {num_workers}")

    process_func = process_task
    
    with Pool(num_workers) as pool:
        for _ in tqdm(pool.imap_unordered(process_func, task), total=len(task)):
            pass
    exit()
    
    # file_groups = file_groups[:5]
    # run
    for i, group in enumerate(tqdm(task)):
        process_func(group)
        exit()

